import json
import math

import numpy as np
import networkx as nx
from switchesalgo import allocate_switches_with_es
from topo1 import build_hybrid_topology_shared_tor
from find_min_max_delay_path import optimize_reconfiguration


# ------------------ 核心函数封装 -------------------
# def load_traffic_data(filename, target_timestamps):
#     """加载指定时间戳的流量数据"""
#     with open(filename, 'r') as f:
#         data = json.load(f)
#
#     traffic_matrices = {}
#     for record in data:
#         ts = record['timestamp']
#         if ts in target_timestamps:
#             traffic_matrices[ts] = {
#                 'matrix': np.array(record['matrix']),
#                 'type': record['type']
#             }
#     return traffic_matrices

# 修改
# def load_traffic_data(filename, target_timestamps=None,traffic_type):
#     """加载流量数据，支持全量加载"""
#     with open(filename, 'r') as f:
#         data = json.load(f)
#
#     traffic_data = {}
#     for record in data:
#         ts = record['timestamp']
#         # 全量加载逻辑：当 target_timestamps=None 时加载所有数据
#         if (target_timestamps is None) or (ts in target_timestamps):
#             traffic_data[ts] = {
#                 'matrix': np.array(record['matrix']),
#                 'type': record['type']
#             }
#     return traffic_data

def load_traffic_data(filename, target_timestamps=None, traffic_type=None):
    """
    加载网络流量矩阵数据，支持多维度过滤

    参数:
        filename (str): JSON文件路径
        target_timestamps (list/None): 需加载的时间戳列表，None表示全量加载
        traffic_type (str/None): 过滤的流量类型：
            - "all_to_all+base" (蓝色曲线基准流量)
            - "all_to_all+burst" (红色曲线突发流量)
            - "mixed" (绿色曲线混合状态)
            - None表示加载全部类型

    返回:
        dict: {timestamp: {'matrix': ndarray, 'type': str}} 的结构化数据
    """
    VALID_TYPES = {"all_to_all+base", "all_to_all+burst", "ring_allreduce+base", "ring_allreduce+burst"}

    with open(filename, 'r') as f:
        raw_data = json.load(f)

    filtered_data = {}
    for record in raw_data:
        # 时间戳过滤
        ts = record['timestamp']
        if target_timestamps and ts not in target_timestamps:
            continue

        # 流量类型验证
        record_type = record['type']
        if traffic_type and record_type != traffic_type:
            continue

        # 类型有效性检查
        if record_type not in VALID_TYPES:
            raise ValueError(f"异常流量类型: {record_type}. 允许类型: {VALID_TYPES}")

        # 矩阵维度转换
        matrix = np.array(record['matrix'])
        if matrix.shape != (16, 16):
            raise ValueError(f"矩阵维度错误，应为16×16，实际是{matrix.shape}")

        filtered_data[ts] = {
            'matrix': matrix,
            'type': record_type,
            'total_traffic': np.sum(matrix)  # 增加总流量预处理
        }

    return filtered_data


def matrix_to_communication_pairs(matrix, node_prefix="ToR"):
    """将流量矩阵转换为通信对字典"""
    # 转换为NumPy数组
    if not isinstance(matrix, np.ndarray):
        matrix = np.array(matrix)
    pairs_dict = {}
    rows, cols = matrix.shape
    for i in range(rows):
        for j in range(cols):
            src = f"{node_prefix}{i + 1}"
            dst = f"{node_prefix}{j + 1}"
            pairs_dict[(src, dst)] = float(matrix[i][j])
    return pairs_dict


def calculate_path_delays(path_dict, traffic_dict, topology):
    """计算路径集合的最大延迟"""
    delays = {}
    for pair, path in path_dict.items():
        # 忽略没有路径的通信对
        if not path:
            continue
        # 计算路径中的最小容量
        min_capacity = float('inf')
        # 遍历路径中的每条链路
        for i in range(len(path) - 1):
            u, v = path[i], path[i + 1]
            if topology.has_edge(u, v):
                capacity = topology[u][v]['capacity']
                min_capacity = min(min_capacity, capacity)

        if min_capacity > 0:
            delays[pair] = traffic_dict[pair] / min_capacity
    return delays


def analyze_reconfiguration(path_dict, topology):
    """分析需要重构的电层路径"""
    electric_pairs = []
    for pair, path in path_dict.items():
        # 忽略没有路径的通信对
        if not path:
            continue

        all_electric = True
        for i in range(len(path) - 1):
            u, v = path[i], path[i + 1]
            if topology[u][v]['type'] != 'electrical':
                all_electric = False
                break

        if all_electric:
            electric_pairs.append(pair)
    return electric_pairs


# ------------------ 业务逻辑函数 -------------------
def process_traffic_scenario(topology, traffic_matrix, switch_params):
    """处理单个流量场景"""
    # 转换为通信对
    traffic_dict = matrix_to_communication_pairs(traffic_matrix)
    pairs = [pair for pair, demand in traffic_dict.items() if demand > 0]

    # 计算最优路径
    switches, paths = allocate_switches_with_es(
        topology, pairs, traffic_dict,
        switch_params['tor_nodes'],
        switch_params['os_nodes'],
        switch_params['threshold']
    )
    # print("-----------基础流量最优测试结果-------------")
    # for pair, path in paths.items():
    #     print(f"通信对: {pair}, 路径: {path}")
    # # print(f"基础流量路径分配结果: {paths}")
    # print("-----------测试结果-------------")

    # 返回前过滤无效路径
    valid_paths = {}
    for pair, path in paths.items():
        if path and isinstance(path, list) and len(path) >= 2:
            valid_paths[pair] = path
        else:
            print(f"丢弃无效路径: {pair}")
    return paths, traffic_dict


print("=======================光电链路重构算法测试===========")
def calculate_total_delay(path_t, reconfigured_paths, no_delays, delays, hybrid_float_values,
                          optical_reconstruction_delay=1, optical_capacity=200, verbose=True):
    """
    计算网络重构后的总完成时延

    参数:
        path_t (dict): 基础路径字典 {通信对: 路径}
        reconfigured_paths (dict): 需要重构的路径字典 {通信对: 路径}
        no_delays (dict): 不重构通信对的时延 {通信对: 时延}
        delays (dict): 重构路径时延 {通信对: 时延}
        hybrid_float_values (dict): 流量需求矩阵 {通信对: 流量需求}
        optical_reconstruction_delay (int): 光链路重构时延(ms)，默认1
        optical_capacity (int): 光链路容量(Gbps)，默认200
        verbose (bool): 是否输出调试信息，默认True

    返回:
        float: 系统总完成时延(ms)
    """

    # ---------------------- 不重构通信对处理 ----------------------
    max_no_pair_delay = 0
    for pair in path_t:
        if pair not in reconfigured_paths:
            current_delay = no_delays.get(pair, 0)
            if current_delay > max_no_pair_delay:
                max_no_pair_delay = current_delay

    if verbose:
        print(f"[INFO] 不需要重构的通信对最大时延: {max_no_pair_delay:.2f} ms")

    # ---------------------- 重构通信对处理 ----------------------
    max_pair_delay = 0

    for pair, path in reconfigured_paths.items():
        if path is None:
            if verbose:
                print(f"[WARN] 通信对 {pair} 无可用路径")
            continue

        pair_delay = delays.get(pair, 0)

        # 处理无穷大时延（电链路容量不足）
        if math.isinf(pair_delay):
            if verbose:
                print(f"[ERROR] 通信对 {pair} 时延为无穷大，请检查链路容量")
            continue

        # 时延优化逻辑
        if pair_delay > optical_reconstruction_delay:
            # 拆分流量：电网络传输部分 + 光网络传输剩余部分
            min_link_capacity = hybrid_float_values[pair] / pair_delay
            remaining_traffic = hybrid_float_values[pair] - min_link_capacity * optical_reconstruction_delay

            if remaining_traffic > 0:
                optical_time = remaining_traffic / optical_capacity
                total_delay = optical_reconstruction_delay + optical_time
            else:
                total_delay = optical_reconstruction_delay

            if verbose:
                print(f"[DEBUG] {pair} 拆分传输: "
                      f"电网络传输 {min_link_capacity:.2f} Gbps × {optical_reconstruction_delay}ms, "
                      f"光网络传输剩余 {remaining_traffic:.2f} Gbps,时延为{total_delay}")
        else:
            total_delay = pair_delay

        # 更新最大时延
        if total_delay > max_pair_delay:
            max_pair_delay = total_delay

    if verbose:
        print(f"[INFO] 需要重构的通信对最大时延: {max_pair_delay:.2f} ms")

    # ---------------------- 总时延计算 ----------------------
    total_delay = max(max_no_pair_delay, max_pair_delay)

    if verbose:
        print(f"\n[RESULT] 重构算法系统总完成时延: {total_delay:.2f} ms")

    return total_delay

print("============================半更新路由算法（只考虑路径更新，不考虑重构）=====================")
def semi_calculate_static_latency(path_t, reconfigured_paths, no_delays, delays, hybrid_float_values,
                          optical_capacity=200, verbose=True):
    # ---------------------- 不重构通信对处理 ----------------------
    max_no_pair_delay = 0
    for pair in path_t:
        if pair not in reconfigured_paths:
            current_delay = no_delays.get(pair, 0)
            print(f"不需要重构的通信对：{pair}的延迟为{current_delay}")
            if current_delay > max_no_pair_delay:
                max_no_pair_delay = current_delay
    # print(f"不需要重构的通信对最大时延为{max_no_pair_delay}")

    if verbose:
        print(f"[INFO] 不需要重构的通信对最大时延: {max_no_pair_delay:.2f} ms")

    # ---------------------- 重构通信对处理 ----------------------
    max_pair_delay = 0

    for pair, path in reconfigured_paths.items():
        if path is None:
            if verbose:
                print(f"[WARN] 通信对 {pair} 无可用路径")
            continue

        pair_delay = delays.get(pair, 0)
        print(f"需要重构的通信对：{pair}的延迟为{pair_delay}")

        # 处理无穷大时延（电链路容量不足）
        if math.isinf(pair_delay):
            if verbose:
                print(f"[ERROR] 通信对 {pair} 时延为无穷大，请检查链路容量")
            continue
        if pair_delay >max_no_pair_delay:
            max_no_pair_delay=pair_delay
    print("不需要重构的通信对最大时延为：",max_no_pair_delay)
    # ---------------------- 总时延计算 ----------------------
    total_delay = max(max_no_pair_delay, max_pair_delay)

    if verbose:
        print(f"\n[RESULT] 系统总完成时延: {total_delay:.2f} ms")

    return total_delay

print("==========================半更新算法==================================")





# ------------------ 主程序 -------------------
if __name__ == "__main__":
    # 初始化网络拓扑
    G = build_hybrid_topology_shared_tor(num_tor=4)

    # 加载流量数据
    traffic_data = load_traffic_data(
        'traffic_simulation.json',
        target_timestamps=[0.12, 5.19]
    )

    print("读取的流量数据如下：")
    print(traffic_data)

    # 处理基础流量，获得基础路径和流量字典
    base_paths, base_traffic = process_traffic_scenario(
        G,
        traffic_data[0.12]['matrix'],
        {
            'tor_nodes': {"ToR1", "ToR2", "ToR3", "ToR4"},
            'os_nodes': {"os1", "os2", "os3", "os4"},
            'threshold': 18
        }

    )
    print("=============测试基础路径====================")
    print(f"基础路径{base_paths},长度为{len(base_paths)}")
    print("=============测试基础路径====================")
    print("\n处理基础流量的结果如下：")
    for pair, path in base_paths.items():
        print(f"通信对: {pair}, 路径: {path}")
    # print(f"基础路径: {base_paths}")
    for pair, traffic in base_traffic.items():
        print(f"通信对: {pair}, 流量: {traffic}")
    # print(f"基础流量: {base_traffic}")
    print(f"基础时延: {calculate_path_delays(base_paths, base_traffic, G)}")

    print(f"基础电层路径: {analyze_reconfiguration(base_paths, G)}")

    # 处理混合流量，获得混合路径和流量字典
    hybrid_paths, hybrid_traffic = process_traffic_scenario(
        G,
        traffic_data[5.19]['matrix'],
        {
            'tor_nodes': {"ToR1", "ToR2", "ToR3", "ToR4"},
            'os_nodes': {"os1", "os2", "os3", "os4"},
            'threshold': 40
        }
    )
    print("\n处理混合流量的结果如下：")
    for pair, path in hybrid_paths.items():
        print(f"通信对: {pair}, 路径: {path}")
    # print(f"基础路径: {base_paths}")
    for pair, traffic in hybrid_traffic.items():
        print(f"通信对: {pair}, 流量: {traffic}")
    # print(f"基础流量: {base_traffic}")
    print(f"基础时延: {calculate_path_delays(hybrid_paths, hybrid_traffic, G)}")

    print(f"基础电层路径: {analyze_reconfiguration(hybrid_paths, G)}")

    # 优化重构路径
    recon_paths, recon_delays, static_delays = optimize_reconfiguration(
        hybrid_traffic, G, base_paths, hybrid_paths
    )
    # print("-----------测试重构时间------------------")
    # for recon_delay1 in recon_delays.values():
    #     if recon_delay1 > 0:
    #         print(f"混合时延: {recon_delay1:.2f} ms")
    # print("-----------测试重构时间------------------")

    print("\n优化重构路径的结果如下：")
    print(f"需要重构的通信对: {len(recon_paths)} 条")
    for pair, path in recon_paths.items():
        if path is not None:
            print(f"{pair}: 电网络路径 {path}，最大时延 {recon_delays[pair]:.2f} ms")
        else:
            print(f"通信对: {pair}, 无可用路径")

    # 输出不需要重构的通信对
    print(f"无需重构的通信对：{len(hybrid_paths) - len(recon_paths)}")
    for pair,path in hybrid_paths.items():
        if pair not in recon_paths:
            print(f"{pair}: 网络路径 {path}，最大时延 {static_delays[pair]:.2f} ms")
        else:
            print(f"通信对: {pair}, 无可用路径")

    # 输出最大时延
    # max_delay = 0
    # print(f"\n需要重构系统最大时延: {max((v for v in recon_delays.values() if not math.isinf(v)), default=0):.2f} ms")
    # print(f"\n不需要重构系统最大时延: {max(static_delays.values(), default=0):.2f} ms")
    # if max(static_delays.values(), default=0) > max((v for v in recon_delays.values() if not math.isinf(v)), default=0):
    #     max_delay = max(static_delays.values(), default=0)
    # else :
    #     max_delay = max((v for v in recon_delays.values() if not math.isinf(v)), default=0)
    # print(f"\n系统总的最大时延: {max_delay:.2f} ms")

    max_delay = calculate_total_delay(hybrid_paths, recon_paths, static_delays, recon_delays, hybrid_traffic,
                                  optical_reconstruction_delay=1, optical_capacity=200, verbose=True)
    print(f"\n光电混合更新系统总的最大时延: {max_delay:.2f} ms")


    # 获取半更新算法的时延
    max_semi_delay = semi_calculate_static_latency(hybrid_paths, recon_paths, static_delays, recon_delays, hybrid_traffic,
                                            optical_capacity=200, verbose=True)

    print(f"半更新系统总的最大时延: {max_semi_delay:.2f} ms")



